#!/usr/bin/python3
# -*- coding: utf-8 -*-

"""
# @version: v1.2
# @author : cd
# @Email : 19688513@qq.com
# @Project : new-horizons-engine
# @File : AKShareDataFetcher.py
# @Software: PyCharm
# @time: 2025/5/30 9:45
# @description : AKShare数据获取类，支持自定义延迟时间和重试次数
"""

import re
import urllib3
import logging
import akshare as ak
from typing import Dict, Optional, List, Tuple, Union
import requests
import time
from datetime import datetime, timedelta
import random
import pandas as pd
import numpy as np

# 配置日志
logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
# 禁用不安全的请求警告
urllib3.disable_warnings(urllib3.exceptions.InsecureRequestWarning)


class AKShareDataFetcher:
    """AKShare数据获取类，负责从AKShare获取数据"""

    def __init__(self, max_retries: int = 5, retry_delay: int = 10,
                 request_delay: float = 0.5, proxy_list: List[str] = None):
        """初始化数据获取器"""
        self.max_retries = max_retries
        self.retry_delay = retry_delay
        self.request_delay = request_delay
        self.proxy_list = proxy_list or []
        self.last_request_time = None
        self._query_cache = {}
        self.cache_expiry = timedelta(minutes=30)

    def _request_limiter(self, request_delay: float) -> None:
        """控制请求频率"""
        if self.last_request_time:
            elapsed = (datetime.now() - self.last_request_time).total_seconds()
            if elapsed < request_delay:
                time.sleep(request_delay - elapsed)
        self.last_request_time = datetime.now()

    def _get_proxy(self) -> Optional[dict]:
        """随机获取一个代理"""
        if self.proxy_list:
            proxy = random.choice(self.proxy_list)
            return {"http": proxy, "https": proxy}
        return None

    def safe_akshare_call(self, func, *args,
                          max_retries: int = None,
                          retry_delay: int = None,
                          request_delay: float = None,
                          **kwargs) -> pd.DataFrame:
        """
        安全调用AKShare函数，带重试机制

        Args:
            func: AKShare函数
            *args: 位置参数
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）
            **kwargs: 关键字参数

        Returns:
            pd.DataFrame: 获取的数据
        """
        # 使用自定义参数或默认参数
        max_retries = max_retries if max_retries is not None else self.max_retries
        retry_delay = retry_delay if retry_delay is not None else self.retry_delay
        request_delay = request_delay if request_delay is not None else self.request_delay

        if self.proxy_list:
            proxies = self._get_proxy()
            if proxies:
                kwargs["proxies"] = proxies

        for attempt in range(max_retries):
            try:
                self._request_limiter(request_delay)
                return func(*args, **kwargs)
            except (requests.exceptions.HTTPError,
                    requests.exceptions.ConnectionError,
                    requests.exceptions.Timeout,
                    requests.exceptions.TooManyRedirects) as e:
                logging.warning(f"网络错误 ({e.__class__.__name__})，重试 {attempt + 1}/{max_retries}")
                time.sleep(retry_delay * (attempt + 1))
            except Exception as e:
                logging.error(f"调用 {func.__name__} 失败: {e}")
                if attempt < max_retries - 1:
                    time.sleep(retry_delay * (attempt + 1))
                else:
                    break
        return pd.DataFrame()

    def validate_dates(self, start_date: str, end_date: str) -> bool:
        """
        验证日期格式

        :param start_date: 开始日期 (YYYYMMDD)
        :param end_date: 结束日期 (YYYYMMDD)
        :return: 是否有效
        """
        try:
            datetime.strptime(start_date, "%Y%m%d")
            datetime.strptime(end_date, "%Y%m%d")
            return True
        except ValueError:
            logging.error("日期格式应为 YYYYMMDD")
            return False

    def get_market_prefix(self, symbol: str) -> str:
        """
        根据股票代码获取市场前缀

        :param symbol: 股票代码
        :return: 市场前缀 ('sh', 'sz', 'bj')
        """
        # 如果已经包含前缀，直接返回
        if symbol.startswith(('sh', 'sz', 'bj')):
            return symbol[:2].lower()

        # 清理非数字字符
        clean_symbol = re.sub(r'\D', '', symbol)

        # 根据数字判断市场
        if clean_symbol.startswith(('6', '9', '5')):
            return 'sh'  # 上证
        elif clean_symbol.startswith(('0', '3')):
            return 'sz'  # 深证
        elif clean_symbol.startswith(('4', '8')):
            return 'bj'  # 北证
        else:
            logging.warning(f"无法识别的股票代码: {symbol}, 默认使用''")
            return ''

    def _standardize_output(self, df: pd.DataFrame, source: str) -> pd.DataFrame:
        """
        标准化输出格式和单位，将新浪数据转换为东方财富格式

        :param df: 原始数据DataFrame
        :param source: 数据源标识 ('em', 'sina')
        :return: 标准化后的DataFrame（东方财富格式）
        """
        if df.empty:
            return df

        # 列名标准化映射
        col_maps = {
            'em': {
                '日期': 'date',
                '开盘': 'open',
                '收盘': 'close',
                '最高': 'high',
                '最低': 'low',
                '成交量': 'volume',
                '成交额': 'turnover',
                '振幅': 'amplitude',
                '涨跌幅': 'change_pct',
                '涨跌额': 'change_amt',
                '换手率': 'turnover_rate'
            },
            'sina': {
                'date': 'date',
                'open': 'open',
                'close': 'close',
                'high': 'high',
                'low': 'low',
                'volume': 'volume',
                'amount': 'turnover',  # 新浪的amount对应东方财富的turnover
                'outstanding_share': 'outstanding_share',
                'turnover': 'turnover_rate'  # 新浪的turnover对应东方财富的turnover_rate
            }
        }

        # 应用列名映射
        col_map = col_maps.get(source, {})
        df = df.rename(columns=col_map)

        # 统一处理日期格式
        if 'date' in df.columns:
            df['date'] = pd.to_datetime(df['date'])

        # 确保数值类型
        num_cols = ['open', 'close', 'high', 'low', 'volume', 'turnover',
                    'amplitude', 'change_pct', 'change_amt', 'turnover_rate']
        for col in num_cols:
            if col in df.columns:
                df[col] = pd.to_numeric(df[col], errors='coerce')

        # 对于新浪数据，计算缺失的字段
        if source == 'sina':
            # 确保数据按日期排序（升序）
            df = df.sort_values('date')

            # 计算前一日收盘价
            df['prev_close'] = df['close'].shift(1)

            # 计算涨跌额 = 收盘价 - 前一日收盘价
            if 'change_amt' not in df.columns:
                df['change_amt'] = df['close'] - df['prev_close']

            # 计算涨跌幅 = (收盘价 - 前一日收盘价) / 前一日收盘价 * 100
            if 'change_pct' not in df.columns:
                df['change_pct'] = (df['close'] - df['prev_close']) / df['prev_close'] * 100

            # 计算振幅 = (最高价 - 最低价) / 前一日收盘价 * 100
            if 'amplitude' not in df.columns:
                df['amplitude'] = (df['high'] - df['low']) / df['prev_close'] * 100

            # 计算换手率（如果缺失）
            if 'turnover_rate' not in df.columns:
                if 'outstanding_share' in df.columns:
                    # 换手率 = 成交量(股) / 流动股本(股) * 100
                    df['turnover_rate'] = df['volume'] / df['outstanding_share'] * 100
                else:
                    # 如果缺少流通股本信息，设为NaN
                    df['turnover_rate'] = float('nan')

            # 删除临时列
            if 'prev_close' in df.columns:
                df = df.drop(columns=['prev_close'])

        # 统一字段顺序 - 使用东方财富的字段顺序
        standard_cols = [
            'date', 'open', 'close', 'high', 'low', 'volume', 'turnover',
            'amplitude', 'change_pct', 'change_amt', 'turnover_rate'
        ]

        # 只保留存在的列，并按照标准顺序排列
        available_cols = [col for col in standard_cols if col in df.columns]
        df = df[available_cols]

        # 添加缺失的列（设为NaN）
        for col in standard_cols:
            if col not in df.columns:
                df[col] = float('nan')

        # 重新排序确保所有字段都存在
        df = df[standard_cols]

        # 按日期降序排列（与东方财富一致）
        df = df.sort_values('date', ascending=False)

        return df

    def fetch_em(self, symbol: str, start_date: str, end_date: str,
                 adjust: str = "", period: str = "daily",
                 max_retries: int = 3, retry_delay: int = 5,
                 request_delay: float = 0.5) -> pd.DataFrame:
        """
        获取东方财富历史行情数据

        :param symbol: 股票代码 (带或不带市场标识)
        :param start_date: 开始日期 (YYYYMMDD)
        :param end_date: 结束日期 (YYYYMMDD)
        :param adjust: 复权类型 ('', 'qfq', 'hfq')
        :param period: 周期 ('daily', 'weekly', 'monthly')
        :param max_retries: 最大重试次数
        :param retry_delay: 重试延迟(秒)
        :param request_delay: 请求间隔(秒)
        :return: 标准化历史数据DataFrame
        """
        # 验证日期格式
        if not self.validate_dates(start_date, end_date):
            return pd.DataFrame()

        # 处理股票代码格式 - 东方财富只需要数字部分
        clean_symbol = re.sub(r'\D', '', symbol)[-6:]

        # 缓存键
        cache_key = f"dc_{clean_symbol}_{start_date}_{end_date}_{adjust}_{period}"
        if cache_key in self._query_cache:
            cached_data, timestamp = self._query_cache[cache_key]
            if (datetime.now() - timestamp).total_seconds() < self.cache_expiry.total_seconds():
                return cached_data.copy()

        # 获取数据
        data = self.safe_akshare_call(
            ak.stock_zh_a_hist,
            symbol=clean_symbol,
            period=period,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

        # 标准化并缓存
        if not data.empty:
            data = self._standardize_output(data, 'em')
            self._query_cache[cache_key] = (data.copy(), datetime.now())
        return data

    def fetch_sina(self, symbol: str, start_date: str, end_date: str,
                   adjust: str = "qfq",
                   max_retries: int = 3, retry_delay: int = 5,
                   request_delay: float = 0.5) -> pd.DataFrame:
        """
        获取新浪财经历史行情数据

        :param symbol: 股票代码 (带或不带市场标识)
        :param start_date: 开始日期 (YYYYMMDD)
        :param end_date: 结束日期 (YYYYMMDD)
        :param adjust: 复权类型 ('', 'qfq', 'hfq')
        :param max_retries: 最大重试次数
        :param retry_delay: 重试延迟(秒)
        :param request_delay: 请求间隔(秒)
        :return: 标准化历史数据DataFrame
        """
        # 验证日期格式
        if not self.validate_dates(start_date, end_date):
            return pd.DataFrame()

        # 处理股票代码格式 - 新浪需要带市场前缀
        prefix = self.get_market_prefix(symbol)
        full_symbol = prefix + re.sub(r'\D', '', symbol)[-6:]

        # 缓存键
        cache_key = f"sina_{full_symbol}_{start_date}_{end_date}_{adjust}"
        if cache_key in self._query_cache:
            cached_data, timestamp = self._query_cache[cache_key]
            if (datetime.now() - timestamp).total_seconds() < self.cache_expiry.total_seconds():
                return cached_data.copy()

        # 获取数据
        data = self.safe_akshare_call(
            ak.stock_zh_a_daily,
            symbol=full_symbol,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

        # 标准化并缓存
        if not data.empty:
            data = self._standardize_output(data, 'sina')
            self._query_cache[cache_key] = (data.copy(), datetime.now())
        return data

    def fetch_historical_data_dc(self, symbol: str, start_date: str, end_date: str,
                                 adjust: str = "",
                                 max_retries: int = None,
                                 retry_delay: int = None,
                                 request_delay: float = None) -> pd.DataFrame:
        """
        获取单只股票历史行情数据-东财

        Args:
            symbol (str): 股票代码
            start_date (str): 起始日期 (YYYYMMDD)
            end_date (str): 结束日期 (YYYYMMDD)
            adjust (str): 复权类型 (qfq: 前复权, hfq: 后复权, 空: 不复权)
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）

        Returns:
            pd.DataFrame: 股票历史数据
        """
        try:
            datetime.strptime(start_date, "%Y%m%d")
            datetime.strptime(end_date, "%Y%m%d")
        except ValueError:
            logging.error("日期格式应为 YYYYMMDD")
            return pd.DataFrame()

        cache_key = f"tx_{symbol}_{start_date}_{end_date}_{adjust}"
        if cache_key in self._query_cache:
            cached_data, timestamp = self._query_cache[cache_key]
            if datetime.now() - timestamp < self.cache_expiry:
                return cached_data.copy()

        data = self.safe_akshare_call(
            ak.stock_zh_a_hist,
            symbol=symbol,
            period='daily',
            start_date=start_date,
            end_date=end_date,
            adjust=adjust,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

        if not data.empty:
            # 标准化数据为统一格式
            data = self._standardize_output(data, 'em')
            self._query_cache[cache_key] = (data.copy(), datetime.now())
        return data

    def fetch_historical_data_sina(self, symbol: str, start_date: str, end_date: str,
                                   adjust: str = "qfq",
                                   max_retries: int = None,
                                   retry_delay: int = None,
                                   request_delay: float = None) -> pd.DataFrame:
        """
        获取单只股票每日数据 (使用 stock_zh_a_daily 接口)

        Args:
            symbol (str): 股票代码 (带市场前缀，如 "sh600519" 或 "sz000001")
            start_date (str): 起始日期 (YYYYMMDD)
            end_date (str): 结束日期 (YYYYMMDD)
            adjust (str): 复权类型 (qfq: 前复权, hfq: 后复权, 空: 不复权)
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）

        Returns:
            pd.DataFrame: 股票历史数据
        """
        try:
            datetime.strptime(start_date, "%Y%m%d")
            datetime.strptime(end_date, "%Y%m%d")
        except ValueError:
            logging.error("日期格式应为 YYYYMMDD")
            return pd.DataFrame()

        cache_key = f"daily_{symbol}_{start_date}_{end_date}_{adjust}"
        if cache_key in self._query_cache:
            cached_data, timestamp = self._query_cache[cache_key]
            if datetime.now() - timestamp < self.cache_expiry:
                return cached_data.copy()

        data = self.safe_akshare_call(
            ak.stock_zh_a_daily,
            symbol=symbol,
            start_date=start_date,
            end_date=end_date,
            adjust=adjust,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

        if not data.empty:
            # 标准化数据为统一格式
            data = self._standardize_output(data, 'sina')
            self._query_cache[cache_key] = (data.copy(), datetime.now())
        return data

    def fetch_stock_basic_info(self,
                               max_retries: int = None,
                               retry_delay: int = None,
                               request_delay: float = None) -> pd.DataFrame:
        """
        获取股票基本信息

        Args:
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）

        Returns:
            pd.DataFrame: 股票基本信息
        """
        return self.safe_akshare_call(
            ak.stock_info_a_code_name,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

    def fetch_stock_esg_info(self,
                             max_retries: int = None,
                             retry_delay: int = None,
                             request_delay: float = None) -> pd.DataFrame:
        """
        获取股票ESG信息

        Args:
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）

        Returns:
            pd.DataFrame: 股票ESG信息
        """
        return self.safe_akshare_call(
            ak.stock_esg_zd_sina,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

    def fetch_stock_detail_info(self, symbol: str,
                                max_retries: int = None,
                                retry_delay: int = None,
                                request_delay: float = None) -> pd.DataFrame:
        """
        获取股票详细信息

        Args:
            symbol (str): 股票代码
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）

        Returns:
            pd.DataFrame: 股票详细信息
        """
        return self.safe_akshare_call(
            ak.stock_individual_info_em,
            symbol=symbol,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )

    def fetch_trade_dates(self,
                          max_retries: int = None,
                          retry_delay: int = None,
                          request_delay: float = None) -> pd.DataFrame:
        """
        获取历史交易日期数据

        Args:
            max_retries (int, optional): 自定义最大重试次数
            retry_delay (int, optional): 自定义重试延迟（秒）
            request_delay (float, optional): 自定义请求间隔（秒）

        Returns:
            pd.DataFrame: 交易日期数据
        """
        return self.safe_akshare_call(
            ak.tool_trade_date_hist_sina,
            max_retries=max_retries,
            retry_delay=retry_delay,
            request_delay=request_delay
        )


if __name__ == "__main__":
    # 创建数据获取器实例
    data_fetcher = AKShareDataFetcher(
        max_retries=3,
        retry_delay=5,
        request_delay=1
    )

    # 1. 获取东方财富数据
    print("东方财富数据示例:")
    df_em = data_fetcher.fetch_historical_data_dc(
        symbol="000001",  # 贵州茅台
        start_date="20230101",
        end_date="20230110",
        adjust="qfq"
    )
    print(df_em)

    # 2. 获取新浪数据并转换为东方财富格式
    print("\n新浪数据转换为东方财富格式示例:")
    df_sina = data_fetcher.fetch_historical_data_sina(
        symbol="sz000001",  # 平安银行
        start_date="20230101",
        end_date="20230110",
        adjust="qfq"
    )
    print(df_sina)

    # 1. 获取东方财富数据（上证指数，000001）
    print("获取东方财富数据（上证指数，000001）...")
    df_em = data_fetcher.fetch_em(
        symbol="000001",  # 上证指数
        start_date="20230101",
        end_date="20230110",
        adjust="qfq",  # 前复权
        period="daily"  # 日线数据
    )

    # 显示数据 - 所有字段
    print("\n东方财富数据（前5行，所有字段）：")
    with pd.option_context('display.max_columns', None, 'display.width', None):
        print(df_em.head())

    # 2. 获取新浪财经数据（平安银行，000001）
    print("\n获取新浪财经数据（平安银行，000001）...")
    df_sina = data_fetcher.fetch_sina(
        symbol="sz000001",  # 平安银行（带市场前缀）
        start_date="20230101",
        end_date="20230110",
        adjust="qfq"  # 前复权
    )

    # 显示数据 - 所有字段
    print("\n新浪财经数据（前5行，所有字段）：")
    with pd.option_context('display.max_columns', None, 'display.width', None):
        print(df_sina.head())

    # 3. 比较两个数据源的结构
    print("\n数据字段对比:")
    print("东方财富字段:", df_em.columns.tolist())
    print("新浪转换后字段:", df_sina.columns.tolist())

    # 4. 存储数据（示例）
    # 这里只是打印存储信息，实际应用中可替换为数据库存储或文件存储
    print("\n存储数据示例:")
    print("存储东方财富数据（{}条记录）".format(len(df_em)))
    print("存储新浪转换后数据（{}条记录）".format(len(df_sina)))

    # 实际存储代码可能如下：
    # df_em.to_csv("em_data.csv", index=False)
    # df_sina.to_csv("sina_converted_data.csv", index=False)
    # 或存储到数据库